from openopt import NLP, GLP
from numpy import cos, arange, ones, asarray, abs, zeros, sqrt, asscalar, array
from pylab import legend, show, plot, subplot, xlabel, subplots_adjust
from string import rjust, ljust, expandtabs

import JT
import Calibration

#===============================================================================
#
#    Data Paths
#
#===============================================================================

fixedImageName = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mhd'
#fixedImageName = 'C:/Users/bryan/bryan-code/trunk/Images/CalibratedDRRImage.mhd'
inputVolFileName = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.mha'
staFile = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.sta'
calFileLat = 'C:/Users/bryan/bryan-code/trunk/TestData/ext_cal1.txt'


#===============================================================================
#
#    Setup DRR
#
#===============================================================================

xraycam = JT.XRayCamera()
drr = JT.DRR()
drr.SetXRayCamera(xraycam)
drr.SetVolumeFilename(inputVolFileName)
drr.SetBlackOnWhiteImage()
drr.InteractionOff()

cal = Calibration.ExternalCalibration()
cal.LoadConsolidatedCalibration(calFileLat)
cal.LoadStaFile(staFile)

xraycam.SetExternalCalibration(cal)

vtkMatrix = cal.ConvertNumpyToVTKArray(cal._VolumeT)
drr.SetVolumeUserMatrix(vtkMatrix)

# Set Fixed image using path
fixedImage = JT.FixedImage()
fixedImage.SetFileName(fixedImageName)

reg=JT.NewRegistration()
reg.SetFixedImage(fixedImage.GetImage(0))
reg.SetMovingImageGenerator(drr)
#reg.SetImageMetric('MMI')
#reg.SetImageMetric('NC')
reg.SetImageMetric('MRSD')

#===============================================================================
#
#    Setup Optimizer
#
#===============================================================================
N = 6
M = 5

probName = reg._metricName+'\n'+fixedImageName[fixedImageName.rfind('/')+1:]
#solver='pswarm'
#solver='galileo'
#solver='ralg'
#solver='de'
#solver='scipy_cobyla'

nlp_list=['ralg', 'scipy_cobyla', 'lincher', 'scipy_slsqp', 'ipopt','algencan']
glp_list=['pswarm','galileo','de']

#solver = nlp_list[1]
solver = glp_list[1]

JT.reg_log.writelines((solver,'\n'))

#lowerBounds=array([2,2,2,1,1,5])
lowerBounds=array([2,2,2,1,1,5])
upperBounds=-1*lowerBounds
maxTime=200
maxTimeCPU=200
plotFlag=1

f=reg.GetMetric
x0=array((0,0,3,2.5,-2.5,0))
#x0=array((0,0,0,0,0,0))
#print "Image tweaked: ", x0
#drr.GenerateImage(x0)

#
if solver in nlp_list:
    print "Starting guess: ", x0
    p = NLP(f, x0, lb=lowerBounds,  ub=upperBounds,
        goal = 'minimum', maxIter = 200,  maxFunEvals = 2000,
        maxTime = maxTime,  maxCPUTime = maxTimeCPU, name=probName)
elif solver in glp_list:
    p = GLP(f, lb=lowerBounds,  ub=upperBounds,
        goal = 'minimum', maxIter = 50,  maxFunEvals = 2000,
        maxTime = maxTime,  maxCPUTime = maxTimeCPU, name=probName)
else:
    print "Not familiar with the how to handle %s solver"%solver

#optional: graphic output
#p.plot = 1

r = p.solve(solver, plot=plotFlag, debug=0)
x_opt,  f_opt = r.xf,  r.ff


print "Results:%s x_opt:%s, f_opt:%s"%(str(r),str(x_opt),str(f_opt))
JT.reg_log.close()
